{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "from stable_baselines3 import PPO\n",
    "from stable_baselines3.common.env_util import make_vec_env\n",
    "from stable_baselines3.common.vec_env import VecNormalize, VecVideoRecorder\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[3], line 69\u001b[0m\n\u001b[0;32m     67\u001b[0m \u001b[38;5;66;03m# 训练循环\u001b[39;00m\n\u001b[0;32m     68\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m---> 69\u001b[0m     \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m     70\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconfig\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtotal_timesteps\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m     71\u001b[0m \u001b[43m        \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m     72\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mPPO_\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mupper\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m     73\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\n\u001b[0;32m     74\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     75\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m     76\u001b[0m     \u001b[38;5;66;03m# 保存模型和归一化参数\u001b[39;00m\n\u001b[0;32m     77\u001b[0m     model\u001b[38;5;241m.\u001b[39msave(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhumanoid_walk_ppo\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[1;32mc:\\Users\\xchen\\anaconda3\\envs\\CX_RL\\lib\\site-packages\\stable_baselines3\\ppo\\ppo.py:311\u001b[0m, in \u001b[0;36mPPO.learn\u001b[1;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[0;32m    302\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mlearn\u001b[39m(\n\u001b[0;32m    303\u001b[0m     \u001b[38;5;28mself\u001b[39m: SelfPPO,\n\u001b[0;32m    304\u001b[0m     total_timesteps: \u001b[38;5;28mint\u001b[39m,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    309\u001b[0m     progress_bar: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[0;32m    310\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m SelfPPO:\n\u001b[1;32m--> 311\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlearn\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    312\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtotal_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtotal_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    313\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcallback\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    314\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlog_interval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlog_interval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    315\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtb_log_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtb_log_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    316\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreset_num_timesteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    317\u001b[0m \u001b[43m        \u001b[49m\u001b[43mprogress_bar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprogress_bar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    318\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\xchen\\anaconda3\\envs\\CX_RL\\lib\\site-packages\\stable_baselines3\\common\\on_policy_algorithm.py:324\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.learn\u001b[1;34m(self, total_timesteps, callback, log_interval, tb_log_name, reset_num_timesteps, progress_bar)\u001b[0m\n\u001b[0;32m    321\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39menv \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m    323\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps \u001b[38;5;241m<\u001b[39m total_timesteps:\n\u001b[1;32m--> 324\u001b[0m     continue_training \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollect_rollouts\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrollout_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_rollout_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mn_steps\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    326\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m continue_training:\n\u001b[0;32m    327\u001b[0m         \u001b[38;5;28;01mbreak\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\xchen\\anaconda3\\envs\\CX_RL\\lib\\site-packages\\stable_baselines3\\common\\on_policy_algorithm.py:247\u001b[0m, in \u001b[0;36mOnPolicyAlgorithm.collect_rollouts\u001b[1;34m(self, env, callback, rollout_buffer, n_rollout_steps)\u001b[0m\n\u001b[0;32m    244\u001b[0m             terminal_value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpolicy\u001b[38;5;241m.\u001b[39mpredict_values(terminal_obs)[\u001b[38;5;241m0\u001b[39m]  \u001b[38;5;66;03m# type: ignore[arg-type]\u001b[39;00m\n\u001b[0;32m    245\u001b[0m         rewards[idx] \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgamma \u001b[38;5;241m*\u001b[39m terminal_value\n\u001b[1;32m--> 247\u001b[0m \u001b[43mrollout_buffer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    248\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_last_obs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[0;32m    249\u001b[0m \u001b[43m    \u001b[49m\u001b[43mactions\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    250\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrewards\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    251\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_last_episode_starts\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[0;32m    252\u001b[0m \u001b[43m    \u001b[49m\u001b[43mvalues\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    253\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlog_probs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    254\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    255\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_obs \u001b[38;5;241m=\u001b[39m new_obs  \u001b[38;5;66;03m# type: ignore[assignment]\u001b[39;00m\n\u001b[0;32m    256\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_episode_starts \u001b[38;5;241m=\u001b[39m dones\n",
      "File \u001b[1;32mc:\\Users\\xchen\\anaconda3\\envs\\CX_RL\\lib\\site-packages\\stable_baselines3\\common\\buffers.py:472\u001b[0m, in \u001b[0;36mRolloutBuffer.add\u001b[1;34m(self, obs, action, reward, episode_start, value, log_prob)\u001b[0m\n\u001b[0;32m    469\u001b[0m action \u001b[38;5;241m=\u001b[39m action\u001b[38;5;241m.\u001b[39mreshape((\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn_envs, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maction_dim))\n\u001b[0;32m    471\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mobservations[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpos] \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(obs)\n\u001b[1;32m--> 472\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mactions[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpos] \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43marray\u001b[49m\u001b[43m(\u001b[49m\u001b[43maction\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    473\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrewards[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpos] \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(reward)\n\u001b[0;32m    474\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mepisode_starts[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpos] \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(episode_start)\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# 硬件加速配置\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "torch.backends.cudnn.benchmark = True  # 启用CUDA加速\n",
    "\n",
    "# 超参数配置\n",
    "config = {\n",
    "    \"n_envs\": 2,  # 并行环境数量（建议设置为CPU核心数）\n",
    "    \"total_timesteps\": 5_000_000,  # 总训练步数（至少需要500万步）\n",
    "    \"policy_kwargs\": {\n",
    "        \"net_arch\": {\n",
    "            \"pi\": [512, 512],  # 策略网络结构\n",
    "            \"vf\": [512, 512],  # 值函数网络结构\n",
    "        },\n",
    "        \"activation_fn\": torch.nn.ReLU,\n",
    "    },\n",
    "    \"learning_rate\": 3e-4,\n",
    "    \"batch_size\": 64,  # 匹配GPU显存容量\n",
    "    \"n_steps\": 2048,  # 每环境采样步数\n",
    "    \"gamma\": 0.99,  # 折扣因子\n",
    "    \"gae_lambda\": 0.95,\n",
    "    \"clip_range\": 0.2,\n",
    "    \"ent_coef\": 0.001,  # 适度鼓励探索\n",
    "    \"target_kl\": 0.05,  # KL散度阈值\n",
    "    \"max_grad_norm\": 0.5,  # 梯度裁剪\n",
    "}\n",
    "\n",
    "# 创建并行环境（带观察归一化）\n",
    "env = make_vec_env(\n",
    "    env_id=\"Humanoid-v4\",\n",
    "    n_envs=config[\"n_envs\"],\n",
    "    env_kwargs={\n",
    "        \"render_mode\": None,  # 训练时不渲染\n",
    "        \"exclude_current_positions_from_observation\": False,\n",
    "        \"reset_noise_scale\": 0.0,  # 初始状态噪声\n",
    "    },\n",
    ")\n",
    "env = VecNormalize(env, norm_obs=True, norm_reward=False)\n",
    "\n",
    "# 可选：训练过程视频记录（每10万步保存一次）\n",
    "if config[\"n_envs\"] == 1:  # 视频录制需要单环境\n",
    "    env = VecVideoRecorder(\n",
    "        env,\n",
    "        video_folder=\"./videos/\",\n",
    "        record_video_trigger=lambda x: x % 100000 == 0,\n",
    "        video_length=200,\n",
    "    )\n",
    "\n",
    "# 初始化PPO模型\n",
    "model = PPO(\n",
    "    policy=\"MlpPolicy\",  # 使用内置的 MlpPolicy\n",
    "    env=env,\n",
    "    device=device,\n",
    "    verbose=1,\n",
    "    tensorboard_log=\"./tb_logs/humanoid_walk2/\",\n",
    "    policy_kwargs=config[\"policy_kwargs\"],  # 传递策略网络参数\n",
    "    learning_rate=config[\"learning_rate\"],\n",
    "    batch_size=config[\"batch_size\"],\n",
    "    n_steps=config[\"n_steps\"],\n",
    "    gamma=config[\"gamma\"],\n",
    "    gae_lambda=config[\"gae_lambda\"],\n",
    "    clip_range=config[\"clip_range\"],\n",
    "    ent_coef=config[\"ent_coef\"],\n",
    "    target_kl=config[\"target_kl\"],\n",
    "    max_grad_norm=config[\"max_grad_norm\"],\n",
    ")\n",
    "\n",
    "# 训练循环\n",
    "try:\n",
    "    model.learn(\n",
    "        total_timesteps=config[\"total_timesteps\"],\n",
    "        progress_bar=True,\n",
    "        tb_log_name=f\"PPO_{device.upper()}\",\n",
    "        reset_num_timesteps=False,\n",
    "    )\n",
    "finally:\n",
    "    # 保存模型和归一化参数\n",
    "    model.save(\"humanoid_walk_ppo\")\n",
    "    env.save(\"humanoid_vecnormalize.pkl\")\n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
